import numpy as np
import matplotlib.pyplot as plt

# -------------------------- 1. 准备数据（给定x、y数组，可替换为真实数据）--------------------------
# 模拟真实场景：给定x（特征数组）和y（标签数组），假设y = 2*x + 3 + 随机噪声（模拟真实数据）
np.random.seed(42)  # 固定随机种子，保证结果可复现
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])  # 给定的x数组（9个样本）
true_w = 2.0  # 真实权重（用于生成y，实际场景中未知）
true_b = 3.0  # 真实偏置（用于生成y，实际场景中未知）
y = true_w * x + true_b + np.random.normal(0, 0.5, size=x.shape)  # 给定的y数组（含少量噪声）

print("=" * 50)
print("给定的x、y数组")
print("=" * 50)
print(f"x数组：{x}")
print(f"y数组：{y.round(2)}")  # 保留2位小数显示
print(f"样本数量：{len(x)}")
print("=" * 50)


# -------------------------- 2. 定义核心函数 --------------------------
def predict(x, w, b):
    """
    预测函数：根据当前参数w、b和输入x，计算预测值y_hat
    公式：y_hat = w * x + b
    """
    return w * x + b


def compute_loss(x, y, w, b):
    """
    计算损失函数（均方误差MSE）：衡量预测值与真实值的差距
    公式：Loss = (1/n) * sum( (y_hat - y)^2 )
    """
    n = len(x)
    y_hat = predict(x, w, b)
    loss = np.sum((y_hat - y) ** 2) / n  # 平均损失，避免样本数影响
    return loss


def compute_gradient(x, y, w, b):
    """
    计算梯度：对w和b分别求偏导，得到梯度方向
    返回：dw（loss对w的偏导）、db（loss对b的偏导）
    """
    n = len(x)
    y_hat = predict(x, w, b)
    # 梯度公式推导：d(Loss)/dw = (2/n)*sum( (y_hat - y)*x )
    dw = 2 * np.sum((y_hat - y) * x) / n
    # 梯度公式推导：d(Loss)/db = (2/n)*sum( (y_hat - y) )
    db = 2 * np.sum(y_hat - y) / n
    return dw, db


# -------------------------- 3. 梯度下降迭代优化 --------------------------
def gradient_descent_fit(
        x, y,  # 给定的x、y数组
        init_w,  # w的初始值
        init_b,  # b的初始值
        learning_rate,  # 学习率（步长）
        iterations,  # 最大迭代次数
        tol=1e-6  # 收敛阈值（损失变化量<tol时停止）
):
    """
    基于梯度下降优化线性模型参数w和b
    返回：迭代过程中的w、b、损失值，以及最终参数
    """
    w = init_w  # 初始化权重
    b = init_b  # 初始化偏置
    loss_history = []  # 记录每次迭代的损失（用于可视化收敛过程）
    w_history = []  # 记录每次迭代的w（用于可视化参数变化）
    b_history = []  # 记录每次迭代的b（用于可视化参数变化）

    for i in range(iterations):
        # 1. 计算当前参数的损失
        current_loss = compute_loss(x, y, w, b)
        loss_history.append(current_loss)
        w_history.append(w)
        b_history.append(b)

        # 2. 计算当前参数的梯度（dw和db）
        dw, db = compute_gradient(x, y, w, b)

        # 3. 沿负梯度方向更新参数
        w = w - learning_rate * dw
        b = b - learning_rate * db

        # 4. 收敛判断：如果损失变化量小于阈值，提前停止（避免无效迭代）
        if i > 0 and abs(loss_history[-1] - loss_history[-2]) < tol:
            print(f"迭代 {i + 1} 次后收敛（损失变化量 < {tol}）")
            break

    # 补充最后一次迭代的参数和损失（如果提前收敛）
    final_loss = compute_loss(x, y, w, b)
    loss_history.append(final_loss)
    w_history.append(w)
    b_history.append(b)

    # 输出迭代信息
    if i == iterations - 1:
        print(f"已完成 {iterations} 次迭代（未达到收敛阈值）")

    return (np.array(w_history), np.array(b_history), np.array(loss_history),
            w, b)  # 返回历史记录和最终参数


# -------------------------- 4. 配置超参数并执行优化 --------------------------
# 超参数设置（可调整观察效果）
init_w = 0.0  # w初始值（随便设，比如0）
init_b = 0.0  # b初始值（随便设，比如0）
learning_rate = 0.01  # 学习率（0.001~0.05合适，太大震荡，太小收敛慢）
iterations = 1000  # 最大迭代次数

# 执行梯度下降拟合
w_history, b_history, loss_history, final_w, final_b = gradient_descent_fit(
    x=x, y=y,
    init_w=init_w,
    init_b=init_b,
    learning_rate=learning_rate,
    iterations=iterations
)

# -------------------------- 5. 输出拟合结果 --------------------------
print("\n" + "=" * 50)
print("梯度下降拟合结果总结")
print("=" * 50)
print(f"初始参数：w0 = {init_w}, b0 = {init_b}")
print(f"学习率 η = {learning_rate}")
print(f"总迭代次数 = {len(loss_history) - 1}")  # 因为loss_history多存了最后一次
print(f"最终拟合参数：w = {final_w:.6f}, b = {final_b:.6f}")
print(f"真实参数（模拟用）：w_true = {true_w}, b_true = {true_b}")
print(f"最终损失值（MSE）：{loss_history[-1]:.6f}")
print(f"拟合模型：y_hat = {final_w:.6f}*x + {final_b:.6f}")
print("=" * 50)

# -------------------------- 6. 可视化结果（3个子图全面展示）--------------------------
plt.rcParams['font.sans-serif'] = ['SimHei']  # 支持中文显示
plt.figure(figsize=(15, 10))

# 子图1：原始数据 + 拟合直线（核心结果展示）
plt.subplot(2, 2, 1)
# 绘制原始数据点
plt.scatter(x, y, c='blue', s=80, label='原始数据（x,y）', alpha=0.8)
# 绘制拟合直线（用最终参数w、b）
x_plot = np.linspace(x.min() - 0.5, x.max() + 0.5, 100)  # 扩展x范围，使直线更完整
y_plot = predict(x_plot, final_w, final_b)
plt.plot(x_plot, y_plot, 'r-', linewidth=2, label=f'拟合直线：y = {final_w:.2f}x + {final_b:.2f}')
# 绘制真实直线（仅模拟场景用，真实数据无此步骤）
y_true_plot = predict(x_plot, true_w, true_b)
plt.plot(x_plot, y_true_plot, 'g--', linewidth=2, alpha=0.7, label=f'真实直线：y = {true_w}x + {true_b}')
plt.xlabel('x（特征值）')
plt.ylabel('y（标签值）')
plt.title('原始数据与拟合结果对比')
plt.legend()
plt.grid(True, alpha=0.3)

# 子图2：损失值随迭代次数变化（收敛曲线）
plt.subplot(2, 2, 2)
plt.plot(range(len(loss_history)), loss_history, 'b-', linewidth=2, label='损失值（MSE）')
plt.xlabel('迭代次数')
plt.ylabel('损失值（MSE）')
plt.title('损失值收敛曲线')
plt.legend()
plt.grid(True, alpha=0.3)
plt.yscale('log')  # 对数坐标，更清晰展示损失下降过程

# 子图3：参数w的迭代变化过程
plt.subplot(2, 2, 3)
plt.plot(range(len(w_history)), w_history, 'orange', linewidth=2, label='w的迭代值')
plt.axhline(y=true_w, color='green', linestyle='--', alpha=0.8, label=f'w真实值（{true_w}）')
plt.xlabel('迭代次数')
plt.ylabel('参数w的值')
plt.title('参数w的收敛过程')
plt.legend()
plt.grid(True, alpha=0.3)

# 子图4：参数b的迭代变化过程
plt.subplot(2, 2, 4)
plt.plot(range(len(b_history)), b_history, 'purple', linewidth=2, label='b的迭代值')
plt.axhline(y=true_b, color='green', linestyle='--', alpha=0.8, label=f'b真实值（{true_b}）')
plt.xlabel('迭代次数')
plt.ylabel('参数b的值')
plt.title('参数b的收敛过程')
plt.legend()
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()